feature: DTVB with Swizzling (tensorB) #1562
Merged
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
solaslin
added
the
noCI
solaslin
force-pushed
the
swizzledB
branch
4 times, most recently
from
c2565ef to
84c9d56
Compare
solaslin
added
enhancement
solaslin
requested review from
jichangjichang,
KKyang,
vin-huang,
imcarsonliao,
hcman2,
Serge45,
Jinp800125 and
TonyYHsieh
as code owners
fixed wave_id distribution
solaslin
commented
Feb 4, 2025
Comment on lines
216
to
223
| depthUDiv = kernel["DepthU"] | ||
| # Swizzled for A, TODO- check for SwizzleTensorB | ||
| depthUDiv = "%s%s"%(kernel["DepthU"], "*MI_M") if (tP["isSwizzled"] and tc == 'A') else "%s"%kernel["DepthU"] | ||
| # | ||
| # swizzle | ||
| if (tP["isSwizzled"] and tc == 'A'): | ||
| depthUDiv = "%s%s"%(kernel["DepthU"], "*MI_M") | ||
| elif (tP["isSwizzled"] and tc == 'B'): | ||
| depthUDiv = "%s%s"%(kernel["DepthU"], "*MI_N") | ||
|
|
||
| gsuOffsetStr = "gsuOffset = DepthU*bpeGR*GSUSumIdx" |
There was a problem hiding this comment.
In this PR, all changes about renaming from MI_M to (MI_MN or MI_MorN) is for both A/B
solaslin
commented
Feb 4, 2025
Comment on lines
3168
to
3173
|
|
||
| if tP["isA"] and tP["isSwizzled"]: | ||
| module.addModuleAsFlatItems(self.alignTo("StrideA0I", "StrideA0I", tP["swizzleK"])) | ||
| if tP["isSwizzled"]: | ||
| # "StrideA0I" or "StrideB1J" | ||
| strideName = "Stride%s%s"%(tc,self.states.indexChars[tP["idx"]]) | ||
| module.addModuleAsFlatItems(self.alignTo(strideName, strideName, tP["swizzleK"])) | ||
|
|
There was a problem hiding this comment.
strideName = "Stride%s%s"%(tc,self.states.indexChars[tP["idx"]])
this would be "StrideA0I" or "StrideB1J".
Thank Jimmy for providing this info.
solaslin
commented
Feb 4, 2025
Comment on lines
3414
to
3425
| module.add(SAddU32(sgpr(tmpSgpr), sgpr("SizesSum"), swizzleStrideVal-1)) | ||
| module.add(SLShiftRightB32(dst=sgpr(tmpSgpr), shiftHex=hex(log2(swizzleStrideVal)), src=sgpr(tmpSgpr), comment="SWZ: numKr = DimK / %s"%swizzleStrideVal)) | ||
| module.add(SLShiftRightB32(dst=sgpr(tmpSgpr), shiftHex=hex(log2(swizzleStrideVal)), src=sgpr(tmpSgpr), comment="SWZ-%s: numKr = DimK / %s"%(tc, swizzleStrideVal))) | ||
| WvG_M = kernel["MIWaveGroup"][0] | ||
| module.add(VAndB32(dst=vgpr(qReg), src0=hex(WvG_M-1), src1=vgpr(qReg), comment="SWZ: wave_id (along_M) %= MIWG[0]")) | ||
| module.add(VMulU32U24(dst=vgpr(qReg), src0=sgpr(tmpSgpr), src1=vgpr(qReg), comment="SWZ: wave_id (along_M) *= numKr")) | ||
| if tP["isA"]: | ||
| module.add(VAndB32(dst=vgpr(qReg), src0=hex(WvG_M-1), src1=vgpr(qReg), comment="SWZ-%s: wave_id (along_M) mod MIWG[0]"%tc)) | ||
| module.add(VMulU32U24(dst=vgpr(qReg), src0=sgpr(tmpSgpr), src1=vgpr(qReg), comment="SWZ-%s: wave_id (along_M) *= numKr"%tc)) | ||
| elif tP["isB"]: | ||
| # NB: | ||
| # Calc of w_id is: /= MIWG[0], not %= MIWG[1] | ||
| module.add(VLShiftRightB32(dst=vgpr(qReg), shiftHex=log2(WvG_M), src=vgpr(qReg), comment="SWZ-%s: wave_id (along_N) /= MIWG[0]"%tc)) | ||
| module.add(VMulU32U24(dst=vgpr(qReg), src0=sgpr(tmpSgpr), src1=vgpr(qReg), comment="SWZ-%s: wave_id (along_N) *= numKr"%tc)) | ||
| elif isDTVAB: |
There was a problem hiding this comment.
Might be the most important note:
For swizzling A: the order of wave is wave_id = wave_id % MI_WaveG[0]
For swizzling B: the order of wave is wave_id =wave_id / MI_WaveG[0] (not wave_id % MI_WaveG[1])
solaslin
commented
Feb 4, 2025
Comment on lines
671
to
680
| { | ||
| const auto k = desc.sizes()[0]; | ||
| const auto m = desc.sizes()[1]; | ||
| const auto b = desc.sizes()[2]; | ||
| const auto swizzleK = miK * packK; | ||
| const auto paddedM = (m + miM - 1) / miM * miM; | ||
| const auto paddedK = (k + swizzleK - 1) / swizzleK * swizzleK; | ||
| return paddedM * paddedK * b; | ||
| // TODO: currently [0][1] = k, (m or n) is based on TN, need to make this generic in the future | ||
| const auto k = desc.sizes()[0]; | ||
| const auto m_n = desc.sizes()[1]; | ||
| const auto b = desc.sizes()[2]; | ||
| const auto swizzleK = miK * packK; | ||
| const auto paddedM_N = (m_n + miM_N - 1) / miM_N * miM_N; | ||
| const auto paddedK = (k + swizzleK - 1) / swizzleK * swizzleK; | ||
| return paddedM_N * paddedK * b; | ||
| } |
There was a problem hiding this comment.
All the changes in host code is to make clear that "M" should be "M or N"
solaslin
commented
Feb 4, 2025
Comment on lines
+21
to
+28
|
|
||
| BenchmarkProblems: | ||
| ######################################## | ||
| # HHS TN DTVB + SWIZZLED_B + BIAS + Activation + SAV | ||
| ######################################## | ||
| - | ||
| - # ProblemType | ||
| OperationType: GEMM |
There was a problem hiding this comment.
Added similar tests for SwizzleB from SwizzleA
mengzcai
approved these changes
Feb 7, 2025
geotseng-amd
approved these changes
Feb 7, 2025
StaceyLai
approved these changes
Feb 7, 2025
|
The Failures in CI are the same as other PSs, which are "Solution not found" issues and have nothing to do with this PR. Merging. |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Resolved SWDEV-509997
[gw2] [ 25%] PASSED Tensile/Tests/common/test_config.py::test_config[Tensile/Tests/self_test/swizzleB.yaml]
[gw3] [ 50%] PASSED Tensile/Tests/common/test_config.py::test_config[Tensile/Tests/self_test/swizzleA.yaml]
[gw0] [ 75%] PASSED Tensile/Tests/common/test_config.py::test_config[Tensile/Tests/self_test/dtl.yaml]
[gw1] [100%] PASSED Tensile/Tests/common/test_config.py::test_config[Tensile/Tests/self_test/dtv.yaml]
=== 4 passed, 5 warnings in 2335.54s (0:38:55) ===
py310: OK (2487.54=setup[10.66]+cmd[0.39,140.81,2335.68] seconds)
congratulations :) (2487.61 seconds)